import { warn } from 'cc';
import Macro from '../../Macro';
import { ActionSpec } from '../Actuators/ActionSpec';
import { ActionBuffers } from '../Actuators/IActionReceiver';
import { AgentInfo } from '../Agent';
import { ISensor } from '../Sensors/ISensor';
import { ObservationWriter } from '../Sensors/ObservationWriter';
import { SensorShapeValidator } from '../Sensors/SensorShapeValidator';
import { GrpcExtensions } from './GrpcExtensions';
import pb from '../Grpc/CommunicatorObjects/pb.js';
import {
    CCCRLInitParameters,
    CommunicatorInitParameters,
    ICommunicator,
    QuitCommandHandler,
    ResetCommandHandler,
} from './ICommunicator';
import { TrainingAnalytics } from '../Analytics/TrainingAnalytics';
import { SideChannelManager } from '../SideChannels/SideChannelManager';
import { MLAgentsSettingsManager } from '../MLAgentsSettingsManager';
import { ChunkParser, ChunkType } from './ChunkParse';

export class RpcCommunicator implements ICommunicator {
    quitCommandReceived: QuitCommandHandler | null = null!;
    resetCommandReceived: ResetCommandHandler | null = null!;

    private _isOpen: boolean = false;

    private _behaviorNames: string[] = [];

    private _needCommunicateThisStep: boolean = false;

    private _observationWriter: ObservationWriter = new ObservationWriter();

    private _sensorShapeValidator: Map<string, SensorShapeValidator> = new Map();

    private _orderedAgentsRequestingDecisions: Map<string, number[]> = new Map();

    private _currentUnityRlInput: pb.communicator_objects.UnityRLOutputProto =
        new pb.communicator_objects.UnityRLOutputProto();

    private _currentUnityRlOutput: pb.communicator_objects.UnityRLOutputProto =
        new pb.communicator_objects.UnityRLOutputProto();

    private _lasActionsReceived: Map<string, Map<number, ActionBuffers>> = new Map();

    private _sentBrainKeys: string[] = [];

    private _unsentBrainKeys: Map<string, ActionSpec> = new Map();

    private _client: pb.communicator_objects.UnityToExternalProto = null!;

    private _reader: ReadableStreamReader<Uint8Array> = null!;

    static create(): RpcCommunicator {
        return new RpcCommunicator();
    }

    static checkCommunicationVersionsAreCompatible(
        unityCommunicationVersion: string,
        pythonApiVersion: string
    ): boolean {
        return true;
    }

    async initialize(initParameters: CommunicatorInitParameters): Promise<[boolean, CCCRLInitParameters]> {
        if (Macro.MLA_SUPPORTED_TRAINING_PLATFORM) {
            const academyParameters = new pb.communicator_objects.UnityRLInitializationOutputProto();
            academyParameters.packageVersion = initParameters.packageVersion;
            academyParameters.communicationVersion = initParameters.communicationVersion;
            academyParameters.capabilities = GrpcExtensions.capabilitiesToProto(initParameters.cccCapabilities);

            let input: pb.communicator_objects.UnityInputProto | undefined = null!;
            let initializationInput: pb.communicator_objects.UnityInputProto | null | undefined = null;
            try {
                const outputProto = new pb.communicator_objects.UnityOutputProto();
                outputProto.rlInitializationOutput = academyParameters;
                [input, initializationInput] = await this.initializeByPort(initParameters.port, outputProto);
            } catch (error) {
                warn(error);
                this.quitCommandReceived?.();
                return [false, new CCCRLInitParameters()];
            }

            const pythonPackageVersion = initializationInput?.rlInitializationInput?.packageVersion ?? '';
            const pythonCommunicationVersion = initializationInput?.rlInitializationInput?.communicationVersion ?? '';
            TrainingAnalytics.setTrainerInformation(pythonPackageVersion, pythonCommunicationVersion);

            const communicationIsCompatible = RpcCommunicator.checkCommunicationVersionsAreCompatible(
                initParameters.communicationVersion,
                pythonCommunicationVersion
            );

            if (initializationInput && input == null) {
                if (!communicationIsCompatible) {
                    warn('Communicator and Trainer communication are not compatible. ');
                } else {
                    warn('Unknown communication error between Python. ');
                }

                return [false, new CCCRLInitParameters()];
            }

            this.updateEnvironmentWithInput(input?.rlInput as pb.communicator_objects.UnityRLInputProto);
            const initParametersOut = GrpcExtensions.toUnityRLInitParameters(
                initializationInput?.rlInitializationInput as pb.communicator_objects.UnityRLInitializationInputProto
            );
            return [true, initParametersOut];
        } else {
            return [false, new CCCRLInitParameters()];
        }
    }
    frameRequest(bytes: Uint8Array): Uint8Array {
        const frame = new ArrayBuffer(bytes.length + 5);
        new DataView(frame, 1, 4).setUint32(0, bytes.length, false);
        new Uint8Array(frame, 5).set(bytes);
        return new Uint8Array(frame);
    }

    parseChunk(buffer: ArrayBuffer) {
        return new ChunkParser()
            .parse(new Uint8Array(buffer))
            .find(chunk => chunk.chunkType === ChunkType.MESSAGE);
    }

    async initializeByPort(
        port: number,
        unityOutput: pb.communicator_objects.UnityOutputProto
    ): Promise<
        [
            pb.communicator_objects.UnityInputProto | undefined,
            pb.communicator_objects.UnityInputProto | null | undefined
        ]
    > {
        this._isOpen = true;
        const proxyPort = MLAgentsSettingsManager.settings.editorProxyPort;

        const rpcImpl = (method: any, requestData: any, callback: any) => {

            fetch(`http://localhost:${proxyPort}/communicator_objects.UnityToExternalProto/Exchange`, {
                method: 'POST',
                body: this.frameRequest(requestData),
                headers: {
                    'Content-Type': 'application/grpc-web',
                    'x-grpc-web': '1',
                },
            }).then((response) => {

                response.arrayBuffer().then((buffer) => {
                    const chunk = this.parseChunk(buffer);
                    callback(null, chunk && chunk.data ? new Uint8Array(chunk.data) : null);
                });
            }).catch((error) => {
                callback(error, null);
            });


        };

        this._client = pb.communicator_objects.UnityToExternalProto.create(rpcImpl);

        let unityInput: pb.communicator_objects.UnityInputProto = new pb.communicator_objects.UnityInputProto();

        const result = await this._client.exchange(this.wrapMessage(unityOutput, 200));

        const inputMessage = await this._client.exchange(this.wrapMessage(null, 200));

        const unityInput1 = inputMessage?.unityInput as pb.communicator_objects.UnityInputProto;
        if (unityInput1) {
            unityInput = unityInput1;
        }

        if (result?.header?.status !== 200 || inputMessage?.header?.status !== 200) {
            this._isOpen = false;
            this.quitCommandReceived?.();
        }
        return [unityInput, result?.unityInput as pb.communicator_objects.UnityInputProto];
    }

    updateEnvironmentWithInput(rlInput: pb.communicator_objects.UnityRLInputProto): void {
        // SideChannelManager.processSideChannelData(rlInput.sideChannel as Uint8Array);
        this.sendCommandEvent(rlInput.command);
    }

    /**
     *  @en Adds the brain to the list of brains which will be sending information to External.
     *  @zh 将大脑添加到将向外部发送信息的大脑列表中。
     * @param name
     * @param actionSpec
     */
    subscribeBrain(name: string, actionSpec: ActionSpec): void {
        if (this._behaviorNames.includes(name)) {
            return;
        }
        this._behaviorNames.push(name);
        this._currentUnityRlOutput.agentInfos[name] =
            new pb.communicator_objects.UnityRLOutputProto.ListAgentInfoProto();

        this.cacheActionSpec(name, actionSpec);
    }

    updateEnvironmentWidthInput(rlInput: pb.communicator_objects.UnityRLInputProto): void {
        // SideChannelManager.processSideChannelData(rlInput.sideChannel as Uint8Array);
        this.sendCommandEvent(rlInput.command);
    }

    putObservations(brainKey: string, info: AgentInfo, sensors: ISensor[]): void {
        //TODO
        if (!this._sensorShapeValidator.has(brainKey)) {
            this._sensorShapeValidator.set(brainKey, new SensorShapeValidator());
        }
        this._sensorShapeValidator.get(brainKey)?.validateSensors(sensors);

        const agentInfoProto = GrpcExtensions.toAgentInfoProto(info);
        sensors.forEach((sensor) => {
            const obsProto = GrpcExtensions.getObservationProto(sensor, this._observationWriter);
            agentInfoProto.observations.push(obsProto);
        });
        this._currentUnityRlOutput.agentInfos[brainKey].value?.push(agentInfoProto);

        this._needCommunicateThisStep = true;

        if (!this._orderedAgentsRequestingDecisions.has(brainKey)) {
            this._orderedAgentsRequestingDecisions.set(brainKey, []);
        }
        if (!info.done) {
            this._orderedAgentsRequestingDecisions.get(brainKey)?.push(info.episodeId);
        }
        if (!this._lasActionsReceived.has(brainKey)) {
            this._lasActionsReceived.set(brainKey, new Map());
        }

        this._lasActionsReceived.get(brainKey)?.set(info.episodeId, ActionBuffers.empty());
        if (info.done) {
            this._lasActionsReceived.get(brainKey)?.delete(info.episodeId);
        }
    }

    decideBatch(): void {
        if (!this._needCommunicateThisStep) {
            return;
        }

        this._needCommunicateThisStep = false;
        this.sendBatchedMessageHelper();
    }

    wrapMessage(
        content: pb.communicator_objects.UnityOutputProto | null,
        status: number
    ): pb.communicator_objects.UnityMessageProto {
        const header = new pb.communicator_objects.HeaderProto();
        header.status = status;

        const message = new pb.communicator_objects.UnityMessageProto();
        message.header = header;
        if (content !== null) {
            message.unityOutput = content;
        }
        return message;
    }

    sendCommandEvent(command: number): void {
        switch (command) {
            case pb.communicator_objects.CommandProto.QUIT: {
                this.quitCommandReceived?.();
                return;
            }
            case pb.communicator_objects.CommandProto.RESET: {
                const keys = this._orderedAgentsRequestingDecisions.keys();
                for (const key of keys) {
                    this._orderedAgentsRequestingDecisions.set(key, []);
                }
                this.resetCommandReceived?.();
                return;
            }
            default: {
                return;
            }
        }
    }

    cacheActionSpec(behaviorName: string, actionSpec: ActionSpec): void {
        if (this._sentBrainKeys.includes(behaviorName)) {
            return;
        }

        this._unsentBrainKeys.set(behaviorName, actionSpec);
    }

    /**
     * @en Helper method that sends the current UnityRLOutput, receives the next UnityInput and
            Applies the appropriate AgentAction to the agents.

        @zh 帮助方法，发送当前 UnityRLOutput，接收下一个 UnityInput 并将适当的 AgentAction 应用于代理。
     */
    async sendBatchedMessageHelper(): Promise<void> {
        const message = new pb.communicator_objects.UnityOutputProto();
        message.rlOutput = this._currentUnityRlOutput;

        const tempUnityRLInitializationOutPut = this.getTempUnityRlInitializationOutput();
        if (tempUnityRLInitializationOutPut) {
            message.rlInitializationOutput = tempUnityRLInitializationOutPut;
        }

        const messageAggregated = SideChannelManager.getSideChannelMessage();
        message.rlOutput.sideChannel = messageAggregated;

        const input = await this.exchange(message);
        this.updateSentActionSpec(tempUnityRLInitializationOutPut);

        for (const key in this._currentUnityRlOutput.agentInfos) {
            const v = this._currentUnityRlOutput.agentInfos[key];
            v.value = [];
        }

        const rlInput = input?.rlInput;
        if (!rlInput) {
            return;
        }
        const agentActions = rlInput.agentActions;
        if(agentActions && Object.keys(agentActions).length === 0){
            return;
        }
            

        this.updateEnvironmentWidthInput(rlInput as pb.communicator_objects.UnityRLInputProto);

        for (const k in rlInput.agentActions) {
            // const v = rlInput.agentActions[k];
            if (!this._orderedAgentsRequestingDecisions.has(k)) {
                return;
            }

            if (!Reflect.has(rlInput.agentActions, k)) {
                return;
            }

            const agentActions = GrpcExtensions.toAgentActionList(
                rlInput.agentActions[k] as pb.communicator_objects.UnityRLInputProto.ListAgentActionProto
            );
            const numAgents = this._orderedAgentsRequestingDecisions.get(k)?.length ?? 0;
            for (let i = 0; i < numAgents; i++) {
                const agentsAction = agentActions[i];
                const agentId = this._orderedAgentsRequestingDecisions.get(k)?.[i] ?? 0;
                if (this._lasActionsReceived.get(k)?.has(agentId)) {
                    this._lasActionsReceived.get(k)?.set(agentId, agentsAction);
                }
            }
        }

        for (const key of this._orderedAgentsRequestingDecisions.keys()) {
            this._orderedAgentsRequestingDecisions.set(key, []);
        }
    }

    getActions(behaviorName: string, agentId: number): ActionBuffers {
        if (this._lasActionsReceived.has(behaviorName)) {
            const actionBuffers = this._lasActionsReceived.get(behaviorName)?.get(agentId);
            if (actionBuffers) {
                return actionBuffers;
            }
        }
        return ActionBuffers.empty();
    }

    /**
     * @en
     * Send a UnityOutput and receives a UnityInput.
     * @zh
     * 发送 UnityOutput 并接收 UnityInput。
     * @param unityOutput
     */
    async exchange(
        unityOutput: pb.communicator_objects.UnityOutputProto
    ): Promise<pb.communicator_objects.UnityInputProto | null> {
        if (!this._isOpen) {
            return null;
        }

        try {
            const message = await this._client.exchange(this.wrapMessage(unityOutput, 200));

            if (message && message?.header?.status === 200) {
                const input = message.unityInput as pb.communicator_objects.UnityInputProto;
                if (input) {
                    return input;
                }
            }
            this._isOpen = false;
            this.quitCommandReceived?.();
            return null;
        } catch (error) {
            this._isOpen = false;
            this.quitCommandReceived?.();
            return null;
        }
    }

    getTempUnityRlInitializationOutput(): pb.communicator_objects.UnityRLInitializationOutputProto {
        let output: pb.communicator_objects.UnityRLInitializationOutputProto = null!;
        this._unsentBrainKeys.forEach((v, k) => {
            if (Reflect.has(this._currentUnityRlOutput.agentInfos, k)) {
                const valueList = this._currentUnityRlOutput.agentInfos[k].value;
                if (valueList && valueList.length > 0) {
                    if (output == null) {
                        output = new pb.communicator_objects.UnityRLInitializationOutputProto();
                    }
                    const actionSpec = this._unsentBrainKeys.get(k);
                    output.brainParameters.push(
                        GrpcExtensions.toBrainParametersProto(actionSpec as ActionSpec, k, true)
                    );
                }
            }
        });
        return output;
    }

    updateSentActionSpec(output: pb.communicator_objects.UnityRLInitializationOutputProto): void {
        if (!output) {
            return;
        }

        output.brainParameters.forEach((brainParameters) => {
            const key = brainParameters.brainName!;
            this._sentBrainKeys.push(key);
            this._unsentBrainKeys.delete(key);
        });
    }
}
